#include "Precompile.h"
#include "DX11MaterialManager.h"
#include "DX11RenderSystem.h"
#include "DX11VertexShader.h"

namespace Graphics{
	DX11MaterialManager::DX11MaterialManager():m_constantBufferId(0){
	
	}

	DX11MaterialManager::~DX11MaterialManager(){
		std::map<unsigned int, ID3D11Buffer*>::iterator itr1 = m_constantBuffers.begin();

		while(itr1 != m_constantBuffers.end()){
			itr1->second->Release();
			itr1++;
		}
	}

	void DX11MaterialManager::initialize(RenderContext* renderContext){
		m_renderContext = renderContext;
		DX11RenderContext* ctx = (DX11RenderContext*)m_renderContext;

		std::vector<std::string> shaderFiles;
		// Load shaders
		switch(static_cast<DX11RenderSystem*>(ctx->m_renderSystem)->getFeatureLevel()){
		case D3D_FEATURE_LEVEL_9_1:
		case D3D_FEATURE_LEVEL_9_2:
			LOG_ERROR("Unsupported shader model");
			assert(0);
			return;
		case D3D_FEATURE_LEVEL_9_3:
			LOG_ERROR("Unimplemented shader model");
			assert(0);
			shaderFiles = Core::GeneralSettings::getFilesInFolder(std::string("Shaders/ShaderModel_3_0/"));
			break;
		case D3D_FEATURE_LEVEL_10_0:
		case D3D_FEATURE_LEVEL_10_1:
			shaderFiles = Core::GeneralSettings::getFilesInFolder(std::string("Shaders/ShaderModel_4_0/"));
			break;
		case D3D_FEATURE_LEVEL_11_0:
			shaderFiles = Core::GeneralSettings::getFilesInFolder(std::string("Shaders/ShaderModel_5_0/"));
			break;
		}

		for(unsigned int i=0; i<shaderFiles.size(); i++){
			if(shaderFiles[i].find(".hlsl") == std::string::npos)
				continue;

			if(shaderFiles[i].find("PS_") != std::string::npos){
				addPixelShader(GraphicsFactory::createPixelShader(shaderFiles[i]));
			}
			else if(shaderFiles[i].find("VS_") != std::string::npos){
				addVertexShader(GraphicsFactory::createVertexShader(shaderFiles[i]));
			}
		}

		setPixelShader(0);
		setVertexShader(0);

		createConstantBuffers(renderContext);

		// Set samplers
		D3D11_SAMPLER_DESC desc;
		desc.Filter = D3D11_FILTER_ANISOTROPIC;
		desc.AddressU = D3D11_TEXTURE_ADDRESS_WRAP;
		desc.AddressV = D3D11_TEXTURE_ADDRESS_WRAP;
		desc.AddressW = D3D11_TEXTURE_ADDRESS_WRAP;
		desc.MipLODBias = 5;
		desc.MaxAnisotropy = 8;

		ID3D11SamplerState* sampler;
		ctx->m_handle->CreateSamplerState(&desc,&sampler);
		ctx->m_context->PSSetSamplers(0,1,&sampler);
		sampler->Release();

		// Create blend states
		D3D11_BLEND_DESC BlendState;
		ZeroMemory(&BlendState, sizeof(D3D11_BLEND_DESC));
		BlendState.IndependentBlendEnable = FALSE;
		BlendState.AlphaToCoverageEnable = FALSE;

		for(unsigned int i=0; i<8; i++){
			BlendState.RenderTarget[i].BlendEnable = TRUE;

			BlendState.RenderTarget[i].SrcBlend = D3D11_BLEND_SRC_ALPHA;
			BlendState.RenderTarget[i].DestBlend = D3D11_BLEND_INV_SRC_ALPHA;
			BlendState.RenderTarget[i].BlendOp = D3D11_BLEND_OP_ADD;

			BlendState.RenderTarget[i].SrcBlendAlpha = D3D11_BLEND_ONE;
			BlendState.RenderTarget[i].DestBlendAlpha = D3D11_BLEND_ZERO;
			BlendState.RenderTarget[i].BlendOpAlpha = D3D11_BLEND_OP_ADD;

			BlendState.RenderTarget[i].RenderTargetWriteMask = D3D11_COLOR_WRITE_ENABLE_ALL;
		}

		ctx->m_handle->CreateBlendState(&BlendState, &m_blendStateBlend);

		BlendState.IndependentBlendEnable = FALSE;
		BlendState.AlphaToCoverageEnable = FALSE;

		for(unsigned int i=0; i<8; i++)
			BlendState.RenderTarget[0].BlendEnable = FALSE;

		ctx->m_handle->CreateBlendState(&BlendState, &m_blendStateNoBlend);

		float blendFactor[] = {0.5,0.5,0.5,0.5};
		UINT sampleMask   = 0xffffffff;
		ctx->m_context->OMSetBlendState(m_blendStateBlend,blendFactor,sampleMask);

		D3D11_DEPTH_STENCIL_DESC dsDesc;
		// Depth test parameters
		dsDesc.DepthEnable = true;
		dsDesc.DepthWriteMask = D3D11_DEPTH_WRITE_MASK_ALL;
		dsDesc.DepthFunc = D3D11_COMPARISON_LESS_EQUAL;

		// Stencil test parameters
		dsDesc.StencilEnable = true;
		dsDesc.StencilReadMask = 0xFF;
		dsDesc.StencilWriteMask = 0xFF;

		// Stencil operations if pixel is front-facing
		dsDesc.FrontFace.StencilFailOp = D3D11_STENCIL_OP_KEEP;
		dsDesc.FrontFace.StencilDepthFailOp = D3D11_STENCIL_OP_INCR;
		dsDesc.FrontFace.StencilPassOp = D3D11_STENCIL_OP_KEEP;
		dsDesc.FrontFace.StencilFunc = D3D11_COMPARISON_ALWAYS;

		// Stencil operations if pixel is back-facing
		dsDesc.BackFace.StencilFailOp = D3D11_STENCIL_OP_KEEP;
		dsDesc.BackFace.StencilDepthFailOp = D3D11_STENCIL_OP_DECR;
		dsDesc.BackFace.StencilPassOp = D3D11_STENCIL_OP_KEEP;
		dsDesc.BackFace.StencilFunc = D3D11_COMPARISON_ALWAYS;

		// Create depth stencil state
		ctx->m_handle->CreateDepthStencilState(&dsDesc, &m_writeEnabled);
		dsDesc.DepthEnable = false;
		ctx->m_handle->CreateDepthStencilState(&dsDesc, &m_writeDisabled);

		ctx->m_context->OMSetDepthStencilState(m_writeEnabled, 1);
	}

	unsigned int DX11MaterialManager::addConstantBuffer(ID3D11Buffer* buffer){
		unsigned int id = m_constantBufferId++;
		m_constantBuffers[id] = buffer;
		return id;
	}

	void DX11MaterialManager::setMaterial(unsigned int id){
		MaterialManager::setMaterial(id);
		updateMaterialBuffer(id);
	}

	void DX11MaterialManager::setPixelShader(unsigned int id){
		MaterialManager::setPixelShader(id);
	}

	void DX11MaterialManager::setVertexShader(unsigned int id){
		std::map<unsigned int,VertexShaderPtr>::iterator itr = m_vertexShaders.find(id);
		if(itr != m_vertexShaders.end()){
			unsigned int layoutId = boost::static_pointer_cast<DX11VertexShader>(itr->second)->m_inputLayoutId;
			if(m_activeInputLayout != layoutId){
				m_inputLayouts[layoutId]->setLayout(m_renderContext);
				m_activeInputLayout = layoutId;
			}
		}

		MaterialManager::setVertexShader(id);
	}

	void DX11MaterialManager::setComputeShader(unsigned int id){
		MaterialManager::setComputeShader(id);
	}

	void DX11MaterialManager::setDepthStencil(bool enable){
		DX11RenderContext* ctx = (DX11RenderContext*)m_renderContext;
		if(enable)
			ctx->m_context->OMSetDepthStencilState(m_writeEnabled,1);
		else
			ctx->m_context->OMSetDepthStencilState(m_writeDisabled,1);
	}

	void DX11MaterialManager::setInstanceShaderVariables(MeshPtr mesh){
		updateInstanceBuffer(mesh);
	}

	void DX11MaterialManager::setJointFrame(std::vector<Utility::Math::Matrix<4,4>> &joints, int controllerId){
		DX11RenderContext* ctx = (DX11RenderContext*)m_renderContext;

		// Set joints if available
		if(controllerId != -1){
			if(joints.empty()){
				joints = ctx->m_animationManager->getAnimationController(controllerId)->getBindPoseMatrices();
			}

			D3D11_MAPPED_SUBRESOURCE mappedBuffer;
			ctx->m_context->Map(m_constantBuffers[2],0,D3D11_MAP_WRITE_DISCARD,0,&mappedBuffer);
				memcpy(((char*)mappedBuffer.pData),&joints[0],joints.size()*sizeof(Utility::Math::Matrix<4,4>));
			ctx->m_context->Unmap(m_constantBuffers[2],0);
		}
	}

	void DX11MaterialManager::setSharedShaderVariables(ICameraPtr camera){
		updateSharedBuffer(camera);
	}

	void DX11MaterialManager::updateSharedBuffer(ICameraPtr camera){
		DX11RenderContext* ctx = (DX11RenderContext*)m_renderContext;
		
		sharedData data;
		data.camPosition = camera->getPosition();
		data.viewProjMatrix = camera->getViewProjection();
		data.projectionMatrix = camera->getProjectionMatrix();
		data.viewMatrix = camera->getViewMatrix();

		data.sinTime = sin(1.0f);

		// Temp dir light
		data.light.a[0] = 0.0f;
		data.light.a[1] = -0.5f;
		data.light.a[2] = 1.0f;

		D3D11_MAPPED_SUBRESOURCE mappedBuffer;
		ctx->m_context->Map(m_constantBuffers[0],0,D3D11_MAP_WRITE_DISCARD,0,&mappedBuffer);
			memcpy(mappedBuffer.pData,&data,sizeof(data));
		ctx->m_context->Unmap(m_constantBuffers[0],0);
	}
	
	void DX11MaterialManager::updateInstanceBuffer(MeshPtr mesh){
		DX11RenderContext* ctx = (DX11RenderContext*)m_renderContext;
		
		instanceData data;
		data.worldMatrix = mesh->getWorldMatrix();

		// Inverse transpose cheating
		D3DXMATRIX worldMat;
		memcpy(worldMat.m,data.worldMatrix.m_data,sizeof(float)*16);
		float determinant = D3DXMatrixDeterminant(&worldMat);
		D3DXMatrixInverse(&worldMat,&determinant,&worldMat);
		//D3DXMatrixTranspose(&worldMat,&worldMat);

		Utility::Math::Matrix<4,4> tmpMat;
		memcpy(tmpMat.m_data,worldMat.m,sizeof(float)*16);

		data.worldInverseTransposeMatrix = tmpMat;

		D3D11_MAPPED_SUBRESOURCE mappedBuffer;
		ctx->m_context->Map(m_constantBuffers[1],0,D3D11_MAP_WRITE_DISCARD,0,&mappedBuffer);
			memcpy(mappedBuffer.pData,&data,sizeof(data));
		ctx->m_context->Unmap(m_constantBuffers[1],0);
	}

	void DX11MaterialManager::updateMaterialBuffer(unsigned int id){
		DX11RenderContext* ctx = (DX11RenderContext*)m_renderContext;
		
		materialData data;

		Graphics::MaterialPtr material = ctx->m_materialManager->getMaterial(id);
		data.colorAmbient = material->m_colorAmbient;
		data.colorDiffuse = material->m_colorDiffuse;
		data.colorEmissive = material->m_colorEmissive;
		data.colorSpecular = material->m_colorSpecular;
		data.colorTransparent = material->m_colorTransparent;
		data.colorTransparency = material->m_colorTransparency;
		data.colorShininess = material->m_colorShininess;

		// Hax
		float blendFactor[] = {0.5,0.5,0.5,0.5};
		UINT sampleMask   = 0xffffffff;
		if(data.colorTransparent.a[0] > 0.0f){
			ctx->m_context->OMSetBlendState(m_blendStateBlend,blendFactor,sampleMask);
		}
		else{
			ctx->m_context->OMSetBlendState(m_blendStateNoBlend,blendFactor,sampleMask);
		}

		D3D11_MAPPED_SUBRESOURCE mappedBuffer;
		ctx->m_context->Map(m_constantBuffers[3],0,D3D11_MAP_WRITE_DISCARD,0,&mappedBuffer);
			memcpy(mappedBuffer.pData,&data,sizeof(data));
		ctx->m_context->Unmap(m_constantBuffers[3],0);
	}

	void DX11MaterialManager::createConstantBuffers(RenderContext* renderContext){
		// Create constant buffers that will supply shaders with application data.
		DX11RenderContext* ctx = (DX11RenderContext*)renderContext;
		ID3D11Buffer* vs_buffer[3];
		ID3D11Buffer* ps_buffer[1];

		// Fill in a buffer description.
		D3D11_BUFFER_DESC cbDesc;
		cbDesc.ByteWidth = sizeof( sharedData );
		cbDesc.Usage = D3D11_USAGE_DYNAMIC;
		cbDesc.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
		cbDesc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
		cbDesc.MiscFlags = 0;
		cbDesc.StructureByteStride = 0;

		// Vertex shader buffers
		HRESULT hr = ctx->m_handle->CreateBuffer( &cbDesc, NULL, &vs_buffer[0] );

		if( FAILED( hr ) ){
			LOG_ERROR("DX11MaterialManager::createConstantBuffers - Failed in creating constant buffer.");
			assert(0);
		}

		addConstantBuffer(vs_buffer[0]);

		cbDesc.ByteWidth = sizeof( instanceData );
		hr = ctx->m_handle->CreateBuffer( &cbDesc, NULL, &vs_buffer[1] );

		if( FAILED( hr ) ){
			LOG_ERROR("DX11MaterialManager::createConstantBuffers - Failed in creating constant buffer.");
			assert(0);
		}

		addConstantBuffer(vs_buffer[1]);

		cbDesc.ByteWidth = 250*16*sizeof(float);
		hr = ctx->m_handle->CreateBuffer( &cbDesc, NULL, &vs_buffer[2] );

		if( FAILED( hr ) ){
			LOG_ERROR("DX11MaterialManager::createConstantBuffers - Failed in creating constant buffer.");
			assert(0);
		}

		addConstantBuffer(vs_buffer[2]);

		// Set the buffers.
		ctx->m_context->VSSetConstantBuffers( 0, 3, vs_buffer );

		// Pixel shader buffers
		cbDesc.ByteWidth = sizeof( materialData );
		hr = ctx->m_handle->CreateBuffer( &cbDesc, NULL, &ps_buffer[0] );

		if( FAILED( hr ) ){
			LOG_ERROR("DX11MaterialManager::createConstantBuffers - Failed in creating constant buffer.");
			assert(0);
		}

		addConstantBuffer(ps_buffer[0]);

		// Set the buffers.
		ctx->m_context->PSSetConstantBuffers( 0, 1, ps_buffer );
	}
}